package cn.webank.framework.data.jedis.utils;

import java.util.Collection;
import java.util.SortedMap;
import java.util.TreeMap;

import com.google.common.base.Charsets;
import com.google.common.hash.HashFunction;

public class ConsistentHash<T> {

	private final HashFunction hashFunction;
	private final int numberOfReplicas;
	private final SortedMap<Long, T> circle = new TreeMap<Long, T>();

	public ConsistentHash(HashFunction hashFunction, int numberOfReplicas,
			Collection<T> nodes) {

		this.hashFunction = hashFunction;
		this.numberOfReplicas = numberOfReplicas;

		for (T node : nodes) {
			add(node);
		}
	}

	public void add(T node) {
		for (int i = 0; i < numberOfReplicas; i++) {
			circle.put(
					hashFunction
							.hashString(node.toString() + i, Charsets.UTF_8)
							.padToLong(), node);
		}
	}

	public void remove(T node) {
		for (int i = 0; i < numberOfReplicas; i++) {
			circle.remove(hashFunction.hashString(node.toString() + i,
					Charsets.UTF_8).padToLong());
		}
	}

	public T get(String key) {
		if (circle.isEmpty()) {
			return null;
		}

		long hash = hashFunction.hashString(key, Charsets.UTF_8).padToLong();
		if (!circle.containsKey(hash)) {
			//SortedMap<Long, T> tailMap = circle.tailMap(hash);
			//hash = tailMap.isEmpty() ? circle.firstKey() : tailMap.firstKey();
			return null;
		}
		return circle.get(hash);
	}

	public T getNext(String key) {
		if (circle.isEmpty()) {
			return null;
		}

		long hash = hashFunction.hashString(key, Charsets.UTF_8).padToLong() + 1;
		if (!circle.containsKey(hash)) {
			SortedMap<Long, T> tailMap = circle.tailMap(hash);
			hash = tailMap.isEmpty() ? circle.firstKey() : tailMap.firstKey();
		}
		return circle.get(hash);

	}
}
